package com.bootdo.app.common.xss;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.StringWriter;
import java.util.LinkedHashMap;
import java.util.Map;

/**
  * @auther: Along
  * @description: xss 请求过滤包装类
  * @Date: created in 2020/12/8 22:05
  */
public class XSSHttpServletRequestWrapper extends HttpServletRequestWrapper {
    public XSSHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    // 引入HtmlFilter对象，对请求参数进行过滤
    private static final HtmlFilter htmlFilter = new HtmlFilter();
    public String xssEncode(String input){
        return htmlFilter.filter(input==null?"":input);
    }

    /**
     * 过滤请求头
     */
    @Override
    public String getHeader(String name) {
        return xssEncode(super.getHeader(xssEncode(name)));
    }

    /**
     * 过滤请求参数
     */
    @Override
    public String getParameter(String name) {
        return xssEncode(super.getHeader(xssEncode(name)));
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        Map<String,String[]> map = new LinkedHashMap<>();
        for(String key : parameterMap.keySet()){
            String[] strings = parameterMap.get(key);
            for(int i = 0; i < strings.length; i++){
                strings[i]=xssEncode(strings[i]);
            }
            map.put(key,strings);
        }
        return map;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] parameterValues = super.getParameterValues(name);
        for(int i = 0; i < parameterValues.length; i++){
            parameterValues[i] = xssEncode(parameterValues[i]);
        }
        return parameterValues;
    }

    /**
     * 过滤请求二进制字节数据
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        // 如果不是json数据，就不过滤
        if(!"application/json".equalsIgnoreCase(super.getHeader("Content-Type"))){
            return super.getInputStream();
        }
        // 否则过滤json数据
        InputStreamReader inputStreamReader = null;
        StringWriter stringWriter = null;
        try{
            inputStreamReader = new InputStreamReader(super.getInputStream(),"UTF-8");
            char[] buffer = new char[4096];
            stringWriter = new StringWriter();
            int len;
            while ((len = inputStreamReader.read(buffer))!=-1){
                stringWriter.write(buffer,0,len);
            }
            String json = xssEncode(stringWriter.toString());
            final ByteArrayInputStream bis = new ByteArrayInputStream(json.getBytes("UTF-8"));
            return new ServletInputStream() {
                @Override
                public boolean isFinished() {
                    return true;
                }

                @Override
                public boolean isReady() {
                    return true;
                }

                @Override
                public void setReadListener(ReadListener readListener) {
                }

                @Override
                public int read() throws IOException {
                    return bis.read();
                }
            };
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            if(inputStreamReader!=null){
                inputStreamReader.close();
            }
        }
        return super.getInputStream();
    }


}
